#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MCP 测试框架示例

这个文件展示了如何为MCP服务创建完整的测试框架，包括：
- 单元测试
- 集成测试
- 性能测试
- 模拟测试
- 自动化测试套件
"""

import unittest
import asyncio
import json
import tempfile
import os
import time
from unittest.mock import Mock, patch, AsyncMock
from typing import Dict, List, Any, Optional
import subprocess
import threading
from contextlib import contextmanager
from mcp.client.stdio import stdio_client
from mcp.types import Tool, Resource

class MCPTestFramework:
    """MCP测试框架基类"""
    
    def __init__(self, server_command: List[str], server_cwd: Optional[str] = None):
        self.server_command = server_command
        self.server_cwd = server_cwd
        self.server_process = None
        self.client = None
    
    @contextmanager
    def mcp_server(self):
        """MCP服务器上下文管理器"""
        try:
            # 启动服务器
            self.server_process = subprocess.Popen(
                self.server_command,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                cwd=self.server_cwd,
                text=True
            )
            
            # 创建客户端
            self.client = stdio_client(
                self.server_process.stdin,
                self.server_process.stdout
            )
            
            yield self.client
            
        finally:
            # 清理资源
            if self.server_process:
                self.server_process.terminate()
                self.server_process.wait(timeout=5)
            
            if self.client:
                try:
                    asyncio.run(self.client.close())
                except:
                    pass
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """调用工具"""
        if not self.client:
            raise RuntimeError("客户端未初始化")
        
        result = await self.client.call_tool(tool_name, arguments)
        return result
    
    async def list_tools(self) -> List[Tool]:
        """列出所有工具"""
        if not self.client:
            raise RuntimeError("客户端未初始化")
        
        response = await self.client.list_tools()
        return response.tools
    
    async def list_resources(self) -> List[Resource]:
        """列出所有资源"""
        if not self.client:
            raise RuntimeError("客户端未初始化")
        
        response = await self.client.list_resources()
        return response.resources
    
    async def read_resource(self, uri: str) -> Any:
        """读取资源"""
        if not self.client:
            raise RuntimeError("客户端未初始化")
        
        result = await self.client.read_resource(uri)
        return result

class TestGreetingServer(unittest.TestCase):
    """问候服务器测试类"""
    
    def setUp(self):
        """测试设置"""
        self.framework = MCPTestFramework(
            server_command=["python", "main.py"],
            server_cwd=os.path.dirname(os.path.abspath(__file__))
        )
    
    def test_server_startup(self):
        """测试服务器启动"""
        with self.framework.mcp_server() as client:
            # 测试客户端连接
            self.assertIsNotNone(client)
    
    def test_list_tools(self):
        """测试工具列表"""
        async def run_test():
            with self.framework.mcp_server() as client:
                tools = await self.framework.list_tools()
                
                # 验证工具存在
                tool_names = [tool.name for tool in tools]
                self.assertIn("say_hello", tool_names)
                self.assertIn("say_goodbye", tool_names)
                
                # 验证工具描述
                hello_tool = next(tool for tool in tools if tool.name == "say_hello")
                self.assertIsNotNone(hello_tool.description)
        
        asyncio.run(run_test())
    
    def test_say_hello_tool(self):
        """测试say_hello工具"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 测试正常调用
                result = await self.framework.call_tool("say_hello", {"name": "Alice"})
                self.assertIn("Hello, Alice!", str(result))
                
                # 测试空名称
                result = await self.framework.call_tool("say_hello", {"name": ""})
                self.assertIn("Hello, !", str(result))
                
                # 测试特殊字符
                result = await self.framework.call_tool("say_hello", {"name": "张三"})
                self.assertIn("Hello, 张三!", str(result))
        
        asyncio.run(run_test())
    
    def test_say_goodbye_tool(self):
        """测试say_goodbye工具"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 测试正常调用
                result = await self.framework.call_tool("say_goodbye", {"name": "Bob"})
                self.assertIn("Goodbye, Bob!", str(result))
        
        asyncio.run(run_test())
    
    def test_invalid_tool(self):
        """测试无效工具调用"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 测试不存在的工具
                with self.assertRaises(Exception):
                    await self.framework.call_tool("invalid_tool", {})
        
        asyncio.run(run_test())
    
    def test_missing_parameters(self):
        """测试缺少参数"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 测试缺少必需参数
                with self.assertRaises(Exception):
                    await self.framework.call_tool("say_hello", {})
        
        asyncio.run(run_test())
    
    def test_list_resources(self):
        """测试资源列表"""
        async def run_test():
            with self.framework.mcp_server() as client:
                resources = await self.framework.list_resources()
                
                # 验证资源存在
                resource_uris = [resource.uri for resource in resources]
                # 根据实际的资源URI进行验证
                # self.assertTrue(any("greeting://" in uri for uri in resource_uris))
        
        asyncio.run(run_test())

class TestPerformance(unittest.TestCase):
    """性能测试类"""
    
    def setUp(self):
        self.framework = MCPTestFramework(
            server_command=["python", "main.py"]
        )
    
    def test_tool_call_performance(self):
        """测试工具调用性能"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 预热
                await self.framework.call_tool("say_hello", {"name": "test"})
                
                # 性能测试
                start_time = time.time()
                iterations = 100
                
                for i in range(iterations):
                    await self.framework.call_tool("say_hello", {"name": f"user{i}"})
                
                end_time = time.time()
                duration = end_time - start_time
                avg_time = duration / iterations
                
                print(f"\n性能测试结果:")
                print(f"总时间: {duration:.4f}秒")
                print(f"平均时间: {avg_time:.4f}秒/次")
                print(f"吞吐量: {iterations/duration:.2f}次/秒")
                
                # 性能断言（根据实际需求调整）
                self.assertLess(avg_time, 0.1, "平均响应时间应小于100ms")
        
        asyncio.run(run_test())
    
    def test_concurrent_calls(self):
        """测试并发调用"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 并发调用测试
                tasks = []
                concurrent_count = 10
                
                start_time = time.time()
                
                for i in range(concurrent_count):
                    task = self.framework.call_tool("say_hello", {"name": f"user{i}"})
                    tasks.append(task)
                
                results = await asyncio.gather(*tasks)
                
                end_time = time.time()
                duration = end_time - start_time
                
                print(f"\n并发测试结果:")
                print(f"并发数: {concurrent_count}")
                print(f"总时间: {duration:.4f}秒")
                print(f"平均时间: {duration/concurrent_count:.4f}秒/次")
                
                # 验证所有调用都成功
                self.assertEqual(len(results), concurrent_count)
                for i, result in enumerate(results):
                    self.assertIn(f"user{i}", str(result))
        
        asyncio.run(run_test())

class TestMockServer(unittest.TestCase):
    """模拟服务器测试类"""
    
    def setUp(self):
        self.mock_client = Mock()
    
    def test_mock_tool_call(self):
        """测试模拟工具调用"""
        # 设置模拟返回值
        self.mock_client.call_tool.return_value = "Hello, Mock!"
        
        # 调用模拟方法
        result = self.mock_client.call_tool("say_hello", {"name": "Mock"})
        
        # 验证结果
        self.assertEqual(result, "Hello, Mock!")
        self.mock_client.call_tool.assert_called_once_with("say_hello", {"name": "Mock"})
    
    @patch('subprocess.Popen')
    def test_server_startup_mock(self, mock_popen):
        """测试服务器启动模拟"""
        # 设置模拟进程
        mock_process = Mock()
        mock_process.stdin = Mock()
        mock_process.stdout = Mock()
        mock_popen.return_value = mock_process
        
        framework = MCPTestFramework(["python", "main.py"])
        
        with framework.mcp_server():
            # 验证进程被创建
            mock_popen.assert_called_once()
            
            # 验证进程被终止
            mock_process.terminate.assert_called_once()

class TestErrorHandling(unittest.TestCase):
    """错误处理测试类"""
    
    def setUp(self):
        self.framework = MCPTestFramework(
            server_command=["python", "main.py"]
        )
    
    def test_server_crash_recovery(self):
        """测试服务器崩溃恢复"""
        # 这里可以测试服务器异常退出后的处理
        pass
    
    def test_invalid_json_handling(self):
        """测试无效JSON处理"""
        # 测试发送无效JSON时的处理
        pass
    
    def test_timeout_handling(self):
        """测试超时处理"""
        # 测试长时间运行的操作超时处理
        pass

class TestDataValidation(unittest.TestCase):
    """数据验证测试类"""
    
    def setUp(self):
        self.framework = MCPTestFramework(
            server_command=["python", "main.py"]
        )
    
    def test_parameter_validation(self):
        """测试参数验证"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 测试各种无效参数
                test_cases = [
                    {"name": None},  # None值
                    {"name": 123},   # 错误类型
                    {"name": "a" * 1000},  # 过长字符串
                    {"invalid_param": "test"},  # 无效参数名
                ]
                
                for case in test_cases:
                    try:
                        result = await self.framework.call_tool("say_hello", case)
                        # 根据实际行为验证结果
                    except Exception as e:
                        # 验证异常类型和消息
                        pass
        
        asyncio.run(run_test())

class TestIntegration(unittest.TestCase):
    """集成测试类"""
    
    def setUp(self):
        self.framework = MCPTestFramework(
            server_command=["python", "main.py"]
        )
    
    def test_full_workflow(self):
        """测试完整工作流程"""
        async def run_test():
            with self.framework.mcp_server() as client:
                # 1. 列出工具
                tools = await self.framework.list_tools()
                self.assertGreater(len(tools), 0)
                
                # 2. 调用每个工具
                for tool in tools:
                    if tool.name == "say_hello":
                        result = await self.framework.call_tool(tool.name, {"name": "Integration Test"})
                        self.assertIsNotNone(result)
                    elif tool.name == "say_goodbye":
                        result = await self.framework.call_tool(tool.name, {"name": "Integration Test"})
                        self.assertIsNotNone(result)
                
                # 3. 列出资源
                resources = await self.framework.list_resources()
                # 根据实际资源进行验证
        
        asyncio.run(run_test())

class TestSuite:
    """测试套件管理器"""
    
    @staticmethod
    def create_test_suite():
        """创建测试套件"""
        suite = unittest.TestSuite()
        
        # 添加基本功能测试
        suite.addTest(unittest.makeSuite(TestGreetingServer))
        
        # 添加性能测试
        suite.addTest(unittest.makeSuite(TestPerformance))
        
        # 添加模拟测试
        suite.addTest(unittest.makeSuite(TestMockServer))
        
        # 添加错误处理测试
        suite.addTest(unittest.makeSuite(TestErrorHandling))
        
        # 添加数据验证测试
        suite.addTest(unittest.makeSuite(TestDataValidation))
        
        # 添加集成测试
        suite.addTest(unittest.makeSuite(TestIntegration))
        
        return suite
    
    @staticmethod
    def run_tests(verbosity=2):
        """运行所有测试"""
        suite = TestSuite.create_test_suite()
        runner = unittest.TextTestRunner(verbosity=verbosity)
        result = runner.run(suite)
        
        # 输出测试报告
        print(f"\n测试报告:")
        print(f"运行测试数: {result.testsRun}")
        print(f"失败数: {len(result.failures)}")
        print(f"错误数: {len(result.errors)}")
        print(f"跳过数: {len(result.skipped) if hasattr(result, 'skipped') else 0}")
        
        if result.failures:
            print("\n失败的测试:")
            for test, traceback in result.failures:
                print(f"  - {test}: {traceback}")
        
        if result.errors:
            print("\n错误的测试:")
            for test, traceback in result.errors:
                print(f"  - {test}: {traceback}")
        
        return result.wasSuccessful()

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="MCP测试框架")
    parser.add_argument("--test", "-t", help="运行特定测试类")
    parser.add_argument("--verbose", "-v", action="store_true", help="详细输出")
    parser.add_argument("--performance", "-p", action="store_true", help="只运行性能测试")
    
    args = parser.parse_args()
    
    if args.test:
        # 运行特定测试
        suite = unittest.TestLoader().loadTestsFromName(args.test)
        runner = unittest.TextTestRunner(verbosity=2 if args.verbose else 1)
        result = runner.run(suite)
    elif args.performance:
        # 只运行性能测试
        suite = unittest.makeSuite(TestPerformance)
        runner = unittest.TextTestRunner(verbosity=2)
        result = runner.run(suite)
    else:
        # 运行所有测试
        success = TestSuite.run_tests(verbosity=2 if args.verbose else 1)
        if not success:
            exit(1)

if __name__ == "__main__":
    main()